import numpy as np
import time
from pyDOE import lhs
import matplotlib
# matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pickle
import scipy.io
import random

# Setup GPU for training
import torch
import torch.nn as nn
import torch.optim as optim
import os

from pinn_particle_hard import sampling_from_msh
from PPINN.model import NeuralNetwork

os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # CPU:-1; GPU0: 1; GPU1: 0;

random.seed(1234)
np.random.seed(1234)
torch.manual_seed(1234)
layers = np.array([2,30,30,30,30,30,30,30,30,1]) 
class PINN_laminar_flow(nn.Module):
    # Initialize the class
    def __init__(self, X_f, N_u, N_f, device):
        # 还原
        super(PINN_laminar_flow, self).__init__()
        self.x_f = {}
        for key, value in X_f.items():
            self.x_f[key] = X_f[key].clone().detach().requires_grad_(True)
        layers[0] = X_f['interior'].shape[1]
        layers[-1] = 6

        self.fluid_null = torch.zeros((7, N_f), device=device)
        self.BC_null = torch.zeros((2, N_u), device=device)

        self.net = NeuralNetwork(layers, device)
        self.optimizer_LBFGS = torch.optim.LBFGS(self.net.parameters(),
                                           max_iter=100000,
                                           max_eval=100000,
                                           history_size=50,
                                           tolerance_change=1 * np.finfo(float).eps,
                                           line_search_fn="strong_wolfe")
        self.optimizer_Adam = torch.optim.Adam(self.net.parameters(), lr=1e-5)
        self.ls = 0
        self.iter = 0
        self.rho = 1
        self.mu = 0.02
        
        # 接口
        Collo = np.array(self.x_f['interior'].detach().cpu().numpy())
        INLET = np.array(self.x_f['velocity_inlet_0'].detach().cpu().numpy())
        OUTLET = np.array(self.x_f['pressure_outlet_0'].detach().cpu().numpy())
        WALL = np.array(self.x_f['wall_0'].detach().cpu().numpy())
        self.uv_net = self.net
        
        ones_col = np.ones((INLET.shape[0], 1))
        zeros_col = np.zeros((INLET.shape[0], 1))
        INLET = np.hstack((INLET, ones_col, zeros_col))
        lb = [0.0, 0.0]
        ub = [1.0, 0.2]
        
        # 原代码
        # Bounds
        self.lb = torch.tensor(lb, dtype=torch.float32)
        self.ub = torch.tensor(ub, dtype=torch.float32)

        # Mat. properties
        # self.rho = 1.0
        # self.mu = 0.02

        # Collocation point
        self.x_c = torch.tensor(Collo[:, 0:1], dtype=torch.float32, requires_grad=True)
        self.y_c = torch.tensor(Collo[:, 1:2], dtype=torch.float32, requires_grad=True)

        self.x_INLET = torch.tensor(INLET[:, 0:1], dtype=torch.float32, requires_grad=True)
        self.y_INLET = torch.tensor(INLET[:, 1:2], dtype=torch.float32, requires_grad=True)
        self.u_INLET = torch.tensor(INLET[:, 2:3], dtype=torch.float32)
        self.v_INLET = torch.tensor(INLET[:, 3:4], dtype=torch.float32)

        self.x_OUTLET = torch.tensor(OUTLET[:, 0:1], dtype=torch.float32, requires_grad=True)
        self.y_OUTLET = torch.tensor(OUTLET[:, 1:2], dtype=torch.float32, requires_grad=True)

        self.x_WALL = torch.tensor(WALL[:, 0:1], dtype=torch.float32, requires_grad=True)
        self.y_WALL = torch.tensor(WALL[:, 1:2], dtype=torch.float32, requires_grad=True)

        # Define layers
        # self.uv_layers = uv_layers

        self.loss_rec = []

        # Initialize NNs
        # self.uv_net = self.initialize_NN(self.uv_layers)

        # Optimizer for solution
        self.optimizer_Adam = optim.Adam(self.uv_net.parameters(), lr=0.0005)

        # Count for callback function
        self.count = 0

    def initialize_NN(self, layers):
        model = nn.Sequential()
        for i in range(len(layers) - 1):
            model.add_module(f'Linear_{i}', nn.Linear(layers[i], layers[i + 1]))
            if i < len(layers) - 2:
                model.add_module(f'Tanh_{i}', nn.Tanh())
        return model

    def forward(self, x, y):
        psips = self.uv_net(torch.cat([x, y], 1))
        psi = psips[:, 0:1]
        p = psips[:, 1:2]
        s11 = psips[:, 2:3]
        s22 = psips[:, 3:4]
        s12 = psips[:, 4:5]
        u = torch.autograd.grad(psi, y, grad_outputs=torch.ones_like(psi), create_graph=True)[0]
        v = -torch.autograd.grad(psi, x, grad_outputs=torch.ones_like(psi), create_graph=True)[0]
        return u, v, p, s11, s22, s12

    def net_f(self, x, y):
        rho = self.rho
        mu = self.mu
        u, v, p, s11, s22, s12 = self.forward(x, y)

        s11_1 = torch.autograd.grad(s11, x, grad_outputs=torch.ones_like(s11), create_graph=True)[0]
        s12_2 = torch.autograd.grad(s12, y, grad_outputs=torch.ones_like(s12), create_graph=True)[0]
        s22_2 = torch.autograd.grad(s22, y, grad_outputs=torch.ones_like(s22), create_graph=True)[0]
        s12_1 = torch.autograd.grad(s12, x, grad_outputs=torch.ones_like(s12), create_graph=True)[0]

        # Plane stress problem
        u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]
        u_y = torch.autograd.grad(u, y, grad_outputs=torch.ones_like(u), create_graph=True)[0]

        v_x = torch.autograd.grad(v, x, grad_outputs=torch.ones_like(v), create_graph=True)[0]
        v_y = torch.autograd.grad(v, y, grad_outputs=torch.ones_like(v), create_graph=True)[0]

        # f_u:=Sxx_x+Sxy_y
        f_u = rho * (u * u_x + v * u_y) - s11_1 - s12_2
        f_v = rho * (u * v_x + v * v_y) - s12_1 - s22_2

        f_s11 = -p + 2 * mu * u_x - s11
        f_s22 = -p + 2 * mu * v_y - s22
        f_s12 = mu * (u_y + v_x) - s12

        f_p = p + (s11 + s22) / 2

        return f_u, f_v, f_s11, f_s22, f_s12, f_p

    def callback(self, loss):
        self.count = self.count + 1
        self.loss_rec.append(loss)
        print('{} th iterations, Loss: {}'.format(self.count, loss))

    def train(self):
        
        # 接口
        it=10000
        
        # 原代码        
        self.net.train()
        loss_WALL = []
        loss_f = []
        loss_INLET = []
        loss_OUTLET = []

        for it in range(10000):
            self.optimizer_Adam.zero_grad()
            

            def closure():
                self.optimizer_LBFGS.zero_grad()
                f_pred_u, f_pred_v, f_pred_s11, f_pred_s22, f_pred_s12, f_pred_p = self.net_f(self.x_c, self.y_c)
                u_inlet_pred, v_inlet_pred, _, _, _, _ = self.forward(self.x_INLET, self.y_INLET)
                _, _, p_outlet_pred, _, _, _ = self.forward(self.x_OUTLET, self.y_OUTLET)
                u_wall_pred, v_wall_pred, _, _, _, _ = self.forward(self.x_WALL, self.y_WALL)
                
                loss_f = torch.mean(f_pred_u ** 2) + torch.mean(f_pred_v ** 2) + torch.mean(f_pred_s11 ** 2) + torch.mean(
                    f_pred_s22 ** 2) + torch.mean(f_pred_s12 ** 2) + torch.mean(f_pred_p ** 2)
                loss_WALL = torch.mean(u_wall_pred ** 2) + torch.mean(v_wall_pred ** 2)
                loss_INLET = torch.mean((u_inlet_pred - self.u_INLET) ** 2) + torch.mean((v_inlet_pred - self.v_INLET) ** 2)
                loss_OUTLET = torch.mean(p_outlet_pred ** 2)

                loss = loss_f + 2 * (loss_WALL + loss_INLET + loss_OUTLET)
                self.ls = loss

                self.ls.backward()

                self.iter += 1
                if not self.iter % 100:
                    print('Epoch: {0:}, Loss: {1:6.9f}'.format(self.iter, self.ls))

                return self.ls
            
            self.optimizer_Adam.step()
        self.iter = 0
        self.optimizer_LBFGS.step(closure)

        return loss_WALL, loss_INLET, loss_OUTLET, loss_f, self.loss_rec

    def predict(self, x_star, y_star):
        x_star = torch.tensor(x_star, dtype=torch.float32, requires_grad=True)
        y_star = torch.tensor(y_star, dtype=torch.float32, requires_grad=True)
        u_star, v_star, p_star, _, _, _ = self.forward(x_star, y_star)
        return u_star.detach().cpu().numpy(), v_star.detach().cpu().numpy(), p_star.detach().cpu().numpy()

def DelCylPT(XY_c, xc=0.0, yc=0.0, r=0.1):
    '''
    delete points within cylinder
    '''
    dst = np.array([((xy[0] - xc) ** 2 + (xy[1] - yc) ** 2) ** 0.5 for xy in XY_c])
    return XY_c[dst > r, :]


def postProcess(xmin, xmax, ymin, ymax, field_MIXED, s=2, alpha=0.5, marker='o'):
    [x_MIXED, y_MIXED, u_MIXED, v_MIXED, p_MIXED] = field_MIXED

    fig, ax = plt.subplots(nrows=3, ncols=1, figsize=(7, 4))
    fig.subplots_adjust(hspace=0.2, wspace=0.2)

    # Plot MIXED result
    cf = ax[0].scatter(x_MIXED, y_MIXED, c=u_MIXED, alpha=alpha-0.1, edgecolors='none', cmap='rainbow', marker=marker, s=int(s))
    ax[0].axis('square')
    for key, spine in ax[0].spines.items():
        if key in ['right','top','left','bottom']:
            spine.set_visible(False)
    ax[0].set_xticks([])
    ax[0].set_yticks([])
    ax[0].set_xlim([xmin, xmax])
    ax[0].set_ylim([ymin, ymax])
    # cf.cmap.set_under('whitesmoke')
    # cf.cmap.set_over('black')
    ax[0].set_title(r'$u$ (m/s)')
    fig.colorbar(cf, ax=ax[0], fraction=0.046, pad=0.04)

    cf = ax[1].scatter(x_MIXED, y_MIXED, c=v_MIXED, alpha=alpha-0.1, edgecolors='none', cmap='rainbow', marker=marker, s=int(s))
    ax[1].axis('square')
    for key, spine in ax[1].spines.items():
        if key in ['right','top','left','bottom']:
            spine.set_visible(False)
    ax[1].set_xticks([])
    ax[1].set_yticks([])
    ax[1].set_xlim([xmin, xmax])
    ax[1].set_ylim([ymin, ymax])
    # cf.cmap.set_under('whitesmoke')
    # cf.cmap.set_over('black')
    ax[1].set_title(r'$v$ (m/s)')
    fig.colorbar(cf, ax=ax[1], fraction=0.046, pad=0.04)

    cf = ax[2].scatter(x_MIXED, y_MIXED, c=p_MIXED, alpha=alpha, edgecolors='none', cmap='rainbow', marker=marker, s=int(s), vmin=-0.25, vmax=4.0)
    ax[2].axis('square')
    for key, spine in ax[2].spines.items():
        if key in ['right','top','left','bottom']:
            spine.set_visible(False)
    ax[2].set_xticks([])
    ax[2].set_yticks([])
    ax[2].set_xlim([xmin, xmax])
    ax[2].set_ylim([ymin, ymax])
    # cf.cmap.set_under('whitesmoke')
    # cf.cmap.set_over('black')
    ax[2].set_title('Pressure (Pa)')
    fig.colorbar(cf, ax=ax[2], fraction=0.046, pad=0.04)

    plt.savefig('./uvp.png', dpi=300)
    plt.close('all')

if __name__ == "__main__":
    # 还原
    X_f_train = sampling_from_msh('data/PINN2D/mesh.npz', 500, 5000)
    for key, value in X_f_train.items():
        X_f_train[key] = torch.tensor(value, dtype=torch.float32, device='cuda:0')
    N_u=500
    N_f=5000

    with torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'):
        # Train from scratch
        
        model = PINN_laminar_flow(X_f_train, N_u, N_f, 'cuda:0')

    # 新代码
        _, _, _, _, _ = model.train()

        # Get mixed-form PINN prediction
        x_PINN = np.linspace(0, 1.0, 251)
        y_PINN = np.linspace(0, 0.2, 101)
        x_PINN, y_PINN = np.meshgrid(x_PINN, y_PINN)
        x_PINN = x_PINN.flatten()[:, None]
        y_PINN = y_PINN.flatten()[:, None]
        # dst = ((x_PINN - 0.2) ** 2 + (y_PINN - 0.2) ** 2) ** 0.5
        # x_PINN = x_PINN[dst >= 0.05]
        # y_PINN = y_PINN[dst >= 0.05]
        # x_PINN = x_PINN.flatten()[:, None]
        # y_PINN = y_PINN.flatten()[:, None]
        x_PINN = torch.tensor(x_PINN, dtype=torch.float32, device='cuda:0')
        y_PINN = torch.tensor(y_PINN, dtype=torch.float32, device='cuda:0')
        u_PINN, v_PINN, p_PINN = model.predict(x_PINN, y_PINN)
        x_PINN = x_PINN.detach().cpu().numpy()
        y_PINN = y_PINN.detach().cpu().numpy()
        field_MIXED = [x_PINN, y_PINN, u_PINN, v_PINN, p_PINN]
        # Plot the comparison of u, v, p
        postProcess(xmin=0, xmax=1, ymin=0, ymax=0.2, field_MIXED=field_MIXED, s=3, alpha=0.5)

